/*
 * Copyright (C) 2024 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the 'License');
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an 'AS IS' BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef IMAGE_KNIFE_C_KEKE_TRANSFORMATION_MASK_H
#define IMAGE_KNIFE_C_KEKE_TRANSFORMATION_MASK_H

#include "transformation_base.h"
#include "transformation/transformation_option.h"
#include "imageknife_internal.h"
#include "downsampling/down_sampler.h"

namespace ImageKnifePro {
class TransformationMask : public TransformationBase {
public:
    explicit TransformationMask(std::shared_ptr<TransformationOption> option) : option_(option)
    {
        IMAGE_KNIFE_CHECK(option_->imageSrc.IsUndefined());
        CacheKeyGenerator::MemoryKeyParams params = {.imageSource = &option->imageSrc,
                                                     .type = ImageRequestType::MAIN_SRC,
                                                     .signature = option_->signature};
        params.downSamplingInfo = DownSampler::GetDownSamplingInfo({0, 0}, DownSamplingStrategy::DEFAULT);
        auto generator = ImageKnifeInternal::GetInstance().GetCacheKeyGenerator();
        IMAGE_KNIFE_CHECK(generator == nullptr);
        maskMemoryKey_ = generator->GenerateMemoryKey(params);
        maskFileKey_ = generator->GenerateFileKey(&option->imageSrc, option_->signature);
    }

    bool BeforeProcess(std::shared_ptr<ImageKnifeTask> task) override
    {
        if (option_->imageSrc.IsUndefined()) {
            task->EchoError("Mask imageSrc is undefined");
            return false;
        }

        OH_PixelmapNative *maskPixelmap = nullptr;
        if (option_->imageSrc.GetType() == DataType::PIXELMAP) {
            option_->imageSrc.GetPixelMap(maskPixelmap);
        } else {
            maskData_ = ImageKnifeInternal::GetInstance().LoadImageSrcSync(&option_->imageSrc,
                task->GetImageKnifeRequest()->GetImageKnifeOption()->context, maskMemoryKey_, maskFileKey_);
            if (maskData_ == nullptr) {
                task->EchoError("Mask ImageData Empty");
                return false;
            }
            maskPixelmap = maskData_->GetPixelmap();
        }

        if (maskPixelmap == nullptr) {
            task->EchoError("Mask Pixelmap Empty");
            return false;
        }
        maskInfo_ = ParsePixelmapInfoReadPixels(maskPixelmap);
        if (maskInfo_.innerBuffer == nullptr) {
            task->EchoError("Mask Read Pixels Failed");
            return "Read Mask Pixels Failed";
        }
        return true;
    }

    std::string GetTransformInfo() override
    {
        return "Default Mask:" + maskMemoryKey_ + "," + option_->signature;
    }

    virtual std::string Process(OH_PixelmapNative *pixelmapIn, OH_PixelmapNative **pixelmapOut) override
    {
        OH_PixelmapNative *maskPixelmap = nullptr;
        CreatePixelmapWithBuffer(maskInfo_.innerBuffer, maskInfo_, &maskPixelmap);
        IMAGE_KNIFE_CHECK_RETURN(maskPixelmap == nullptr, "Create Mask Pixelmap Failed");

        PixemapInfo info = ParsePixelmapInfoReadPixels(pixelmapIn);
        if (info.innerBuffer == nullptr) {
            return "Read Pixels Failed";
        }

        if (OH_PixelmapNative_Scale(maskPixelmap, info.width / static_cast<float>(maskInfo_.width),
                                    info.height / static_cast<float>(maskInfo_.height)) != IMAGE_SUCCESS) {
            free(info.innerBuffer);
            OH_PixelmapNative_Release(maskPixelmap);
            return "Mask Scale Failed";
        }

        PixemapInfo maskInfoScaled = ParsePixelmapInfoReadPixels(maskPixelmap);
        if (maskInfoScaled.innerBuffer == nullptr) {
            free(info.innerBuffer);
            OH_PixelmapNative_Release(maskPixelmap);
            return "Read Scaled Mask Pixels Failed";
        }
        IMAGE_KNIFE_LOG(LOG_DEBUG, "width %{public}d:%{public}d, height %{public}d:%{public}d",
                        info.width, maskInfoScaled.width, info.height, maskInfoScaled.height);

        Mask(info, maskInfoScaled);

        // 创建新的pixelmap
        CreatePixelmapWithBuffer(info.innerBuffer, info, pixelmapOut);

        free(info.innerBuffer);
        free(maskInfoScaled.innerBuffer);
        OH_PixelmapNative_Release(maskPixelmap);
        return std::string();
    }

    void AfterProcess() override
    {
        if (maskInfo_.innerBuffer != nullptr) {
            free(maskInfo_.innerBuffer);
            maskInfo_.innerBuffer = nullptr;
        }
    }
private:
    std::shared_ptr<TransformationOption> option_;
    std::shared_ptr<ImageData> maskData_ = nullptr;
    PixemapInfo maskInfo_;

    std::string maskMemoryKey_;
    std::string maskFileKey_;

    void Mask(PixemapInfo &info, PixemapInfo &maskInfoScaled)
    {
        const int dimensionality = 4;
        const int blue = 2;
        const int alpha = 3;
        for (int h = 0; h < info.height; h++) {
            for (int w = 0; w < info.width; w++) {
                if (h >= maskInfoScaled.height || w >= maskInfoScaled.width) {
                    continue;
                }

                int maskIndex = (h * maskInfoScaled.width + w) * dimensionality;
                if (maskInfoScaled.innerBuffer[maskIndex + alpha] == 0) {
                    continue;
                }

                int index = (h * info.width + w) * dimensionality;
                info.innerBuffer[index] = maskInfoScaled.innerBuffer[maskIndex];
                info.innerBuffer[index + 1] = maskInfoScaled.innerBuffer[maskIndex + 1];
                info.innerBuffer[index + blue] = maskInfoScaled.innerBuffer[maskIndex + blue];
                info.innerBuffer[index + alpha] = maskInfoScaled.innerBuffer[maskIndex + alpha];
            }
        }
    }
};

}

#endif //IMAGE_KNIFE_C_KEKE_TRANSFORMATION_MASK_H
